#pragma once

#include <opencv2/core.hpp>

template <int num_bins, int factor>
class LocalHist {
public:
	typedef struct {
		float pf;
		float pb;
	} Float2;

	void Clear() {	
		memset(bins, 0, full_hist_size * 2 * sizeof(float));
		etaf = 0, etab = 0;
		we = 1.0f;
	}

	void BuildInRect(cv::Mat& frame, cv::Mat& mask, cv::Rect& rc) {
		float num_fg_pixs = 0, num_bg_pixs = 0;
		etaf = 0, etab = 0;

		for (int r = rc.y; r < rc.y + rc.height; ++r)
		for (int c = rc.x; c < rc.x + rc.width; ++c) {
			uchar* frame_ptr = (uchar*)(frame.data + r * frame.step + 3 * c);
			const int ru = (frame_ptr[2] >> factor);
			const int gu = (frame_ptr[1] >> factor);
			const int bu = (frame_ptr[0] >> factor);
			const int bidx = (ru + gu * num_bins) * num_bins + bu;

			if (mask.at<uchar>(r, c)) {
				bins[bidx].pf += 1;
				num_fg_pixs += 1;
				etaf += 1;
			}	else {
				bins[bidx].pb += 1;
				num_bg_pixs += 1;
				etab += 1;
			}
		}

		num_fg_pixs = (num_fg_pixs != 0) ? 1.f / num_fg_pixs : 0.f;
		num_bg_pixs = (num_bg_pixs != 0) ? 1.f / num_bg_pixs : 0.f;

		for (int i = 0; i < full_hist_size; ++i) {
			bins[i].pf = bins[i].pf * num_fg_pixs;
			bins[i].pb = bins[i].pb * num_bg_pixs;
		}
	}

	void TestLine(uchar* frame_row, uchar* mask_row, int xl, int xr, Float2* bins, float& sum_err, float& sum_all) {
		uchar* frame_ptr = (uchar*)(frame_row)+3 * xl;
		uchar* mask_ptr = (uchar*)(mask_row)+xl;
		uchar* mask_max_ptr = (uchar*)(mask_row)+xr;

		for (; mask_ptr <= mask_max_ptr; mask_ptr += 1, frame_ptr += 3) {
			sum_all += 1;

			int ru = (frame_ptr[2] >> factor);
			int gu = (frame_ptr[1] >> factor);
			int bu = (frame_ptr[0] >> factor);
			//int pidx = (ru * num_bins + gu) * num_bins + bu;
			const int bidx = (ru + gu * num_bins) * num_bins + bu;

			float pf = bins[bidx].pf;
			float pb = bins[bidx].pb;

			pf += 0.0000001f; pb += 0.0000001f;

			float ppf = etaf * pf / (etaf * pf + etab * pb);
			float ppb = etab * pb / (etaf * pf + etab * pb);

			//if (*mask_ptr == _m_id) {
			if ((*mask_ptr != 0 && ppf < ppb) || (*mask_ptr == 0 && ppf > ppb)) {
				sum_err += 1;
			}
		}
	}

	void ProcessLine(uchar* frame_row, uchar* mask_row, int xl, int xr, Float2* bins, float& sum_f, float& sum_b) {
		uchar* frame_ptr = (uchar*)(frame_row)+3 * xl;
		uchar* mask_ptr = (uchar*)(mask_row)+xl;
		uchar* mask_max_ptr = (uchar*)(mask_row)+xr;

		for (; mask_ptr <= mask_max_ptr; mask_ptr += 1, frame_ptr += 3) {
			int ru = (frame_ptr[2] >> factor);
			int gu = (frame_ptr[1] >> factor);
			int bu = (frame_ptr[0] >> factor);
			//int pidx = (ru * num_bins + gu) * num_bins + bu;
			const int pidx = (ru + gu * num_bins) * num_bins + bu;

			//if (*mask_ptr == _m_id) {
			if (*mask_ptr != 0) {
				bins[pidx].pf += 1;
				sum_f += 1;
			}
			else {
				bins[pidx].pb += 1;
				sum_b += 1;
			}
		}
	}

	void ComputeWeight(cv::Mat& frame, cv::Mat& mask, cv::Point& center, int radius) {
		//float num_fg_pixs = 0, num_bg_pixs = 0;
		//etaf = 0, etab = 0;

		uchar* frame_data = frame.data;
		uchar* mask_data = mask.data;

		size_t frame_step = frame.step;
		size_t mask_step = mask.step;

		int err = 0;
		int dx = radius;
		int dy = 0;
		int plus = 1;
		int minus = (radius << 1) - 1;
		int olddx = dx;
		
		int inside = center.x >= radius && center.x < frame.cols - radius && center.y >= radius && center.y < frame.rows - radius;
		float sum_all = 0, sum_err = 0;

		while (dx >= dy) {
			int mask;
			int y11 = center.y - dy, y12 = center.y + dy, y21 = center.y - dx, y22 = center.y + dx;
			int x11 = center.x - dx, x12 = center.x + dx, x21 = center.x - dy, x22 = center.x + dy;

			if (inside) {
				uchar* frame_row = frame_data + y11 * frame_step;
				uchar* mask_row = mask_data + y11 * mask_step;
				TestLine(frame_row, mask_row, x11, x12, bins, sum_err, sum_all);

				if (y11 != y12) {
					uchar* frame_row = frame_data + y12 * frame_step;
					uchar* mask_row = mask_data + y12 * mask_step;
					TestLine(frame_row, mask_row, x11, x12, bins, sum_err, sum_all);
				}

				if (olddx != dx) {
					if (y11 != y21) {
						uchar* frame_row = frame_data + y21 * frame_step;
						uchar* mask_row = mask_data + y21 * mask_step;
						TestLine(frame_row, mask_row, x21, x22, bins, sum_err, sum_all);
					}

					if (y12 != y22) {
						uchar* frame_row = frame_data + y22 * frame_step;
						uchar* mask_row = mask_data + y22 * mask_step;
						TestLine(frame_row, mask_row, x21, x22, bins, sum_err, sum_all);
					}
				}
			}	else if (x11 < frame.cols && x12 >= 0 && y21 < frame.rows && y22 >= 0) {
				x11 = std::max(x11, 0);
				x12 = MIN(x12, frame.cols - 1);

				if ((unsigned)y11 < (unsigned)frame.rows) {
					uchar* frame_row = frame_data + y11 * frame_step;
					uchar* mask_row = mask_data + y11 * mask_step;
					TestLine(frame_row, mask_row, x11, x12, bins, sum_err, sum_all);
				}

				if ((unsigned)y12 < (unsigned)frame.rows && (y11 != y12)) {
					uchar* frame_row = frame_data + y12 * frame_step;
					uchar* mask_row = mask_data + y12 * mask_step;
					TestLine(frame_row, mask_row, x11, x12, bins, sum_err, sum_all);
				}

				if (x21 < frame.cols && x22 >= 0 && (olddx != dx)) {
					x21 = std::max(x21, 0);
					x22 = MIN(x22, frame.cols - 1);

					if ((unsigned)y21 < (unsigned)frame.rows) {
						uchar* frame_row = frame_data + y21 * frame_step;
						uchar* mask_row = mask_data + y21 * mask_step;
						TestLine(frame_row, mask_row, x21, x22, bins, sum_err, sum_all);
					}

					if ((unsigned)y22 < (unsigned)frame.rows) {
						uchar* frame_row = frame_data + y22 * frame_step;
						uchar* mask_row = mask_data + y22 * mask_step;
						TestLine(frame_row, mask_row, x21, x22, bins, sum_err, sum_all);
					}
				}
			}

			olddx = dx;

			dy++;
			err += plus;
			plus += 2;

			mask = (err <= 0) - 1;

			err -= minus & mask;
			dx += mask;
			minus -= mask & 2;
		}

		we = (sum_all - sum_err) / sum_all;
	}

	void BuildInCircle(cv::Mat& frame, cv::Mat& mask, cv::Point& center, int radius) {
		//float num_fg_pixs = 0, num_bg_pixs = 0;
		etaf = 0, etab = 0;

		uchar* frame_data = frame.data;
		uchar* mask_data = mask.data;

		size_t frame_step = frame.step;
		size_t mask_step = mask.step;

		int err = 0;
		int dx = radius;
		int dy = 0;
		int plus = 1;
		int minus = (radius << 1) - 1;
		int olddx = dx;
		
		int inside = center.x >= radius && center.x < frame.cols - radius && center.y >= radius && center.y < frame.rows - radius;
	
		while (dx >= dy) {
			int mask;
			int y11 = center.y - dy, y12 = center.y + dy, y21 = center.y - dx, y22 = center.y + dx;
			int x11 = center.x - dx, x12 = center.x + dx, x21 = center.x - dy, x22 = center.x + dy;

			if (inside) {
				uchar* frame_row = frame_data + y11 * frame_step;
				uchar* mask_row = mask_data + y11 * mask_step;
				ProcessLine(frame_row, mask_row, x11, x12, bins, etaf, etab);

				if (y11 != y12) {
					uchar* frame_row = frame_data + y12 * frame_step;
					uchar* mask_row = mask_data + y12 * mask_step;
					ProcessLine(frame_row, mask_row, x11, x12, bins, etaf, etab);
				}

				if (olddx != dx) {
					if (y11 != y21) {
						uchar* frame_row = frame_data + y21 * frame_step;
						uchar* mask_row = mask_data + y21 * mask_step;
						ProcessLine(frame_row, mask_row, x21, x22, bins, etaf, etab);
					}

					if (y12 != y22) {
						uchar* frame_row = frame_data + y22 * frame_step;
						uchar* mask_row = mask_data + y22 * mask_step;
						ProcessLine(frame_row, mask_row, x21, x22, bins, etaf, etab);
					}
				}
			}	else if (x11 < frame.cols && x12 >= 0 && y21 < frame.rows && y22 >= 0) {
				x11 = std::max(x11, 0);
				x12 = MIN(x12, frame.cols - 1);

				if ((unsigned)y11 < (unsigned)frame.rows) {
					uchar* frame_row = frame_data + y11 * frame_step;
					uchar* mask_row = mask_data + y11 * mask_step;
					ProcessLine(frame_row, mask_row, x11, x12, bins, etaf, etab);
				}

				if ((unsigned)y12 < (unsigned)frame.rows && (y11 != y12)) {
					uchar* frame_row = frame_data + y12 * frame_step;
					uchar* mask_row = mask_data + y12 * mask_step;
					ProcessLine(frame_row, mask_row, x11, x12, bins, etaf, etab);
				}

				if (x21 < frame.cols && x22 >= 0 && (olddx != dx)) {
					x21 = std::max(x21, 0);
					x22 = MIN(x22, frame.cols - 1);

					if ((unsigned)y21 < (unsigned)frame.rows) {
						uchar* frame_row = frame_data + y21 * frame_step;
						uchar* mask_row = mask_data + y21 * mask_step;
						ProcessLine(frame_row, mask_row, x21, x22, bins, etaf, etab);
					}

					if ((unsigned)y22 < (unsigned)frame.rows) {
						uchar* frame_row = frame_data + y22 * frame_step;
						uchar* mask_row = mask_data + y22 * mask_step;
						ProcessLine(frame_row, mask_row, x21, x22, bins, etaf, etab);
					}
				}
			}

			olddx = dx;

			dy++;
			err += plus;
			plus += 2;

			mask = (err <= 0) - 1;

			err -= minus & mask;
			dx += mask;
			minus -= mask & 2;
		}

		float num_fg_pixs = (etaf != 0) ? 1.f / etaf : 0.f;
		float num_bg_pixs = (etab != 0) ? 1.f / etab : 0.f;

		for (int i = 0; i < full_hist_size; ++i) {
			bins[i].pf = bins[i].pf * num_fg_pixs;
			bins[i].pb = bins[i].pb * num_bg_pixs;
		}
	}

	void Merge(const LocalHist<num_bins, factor>& hist, float afg, float abg) {
		for (int i = 0; i < full_hist_size; ++i) {
			bins[i].pf = bins[i].pf * afg + hist.bins[i].pf * (1.0f - afg);
			bins[i].pb = bins[i].pb * abg + hist.bins[i].pb * (1.0f - abg);
		}
	}

	void GetPosterProb(unsigned char r, unsigned char g, unsigned char b, float& ppf, float& ppb) {
		const int ru = (r >> factor);
		const int gu = (g >> factor);
		const int bu = (b >> factor);
		const int bidx = (ru + gu * num_bins) * num_bins + bu;

		float pf = bins[bidx].pf;
		float pb = bins[bidx].pb;

		pf += 0.0000001f; pb += 0.0000001f;

		ppf = etaf * pf / (etaf * pf + etab * pb);
		ppb = etab * pb / (etaf * pf + etab * pb);
		
		//ppf = pf / (pf + pb);
		//ppb = pb / (pf + pb);
	}

	static const int full_hist_size = num_bins * num_bins * num_bins;
	Float2 bins[num_bins*num_bins*num_bins];
	float etaf, etab;
	float we;
	//int num_fg_pixs, num_bg_pixs;
};

typedef LocalHist<64,2> LocalHist64;
typedef LocalHist<32,3> LocalHist32;
typedef LocalHist<16,4> LocalHist16;